# -*- coding: utf-8 -*-
"""
@Author: Ling
@Date: 2023/3/11 23:45 
"""
# -*- coding: utf-8 -*-
import pandas as pd
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import os
import torch
import cv2
import albumentations as A
from albumentations.pytorch import ToTensorV2


class MyTestDataSet(Dataset):
    def __init__(self, dir,csv_path):
        df=pd.read_csv(csv_path)
        train_dic = df.iloc[:, 1].unique()
        self.train_to_dic = {j: i for i, j in enumerate(train_dic)}  # 从文字转为数字标签
        self.train_from_dic = {i: j for i, j in enumerate(train_dic)}  # 从数字转为文字标签
        self.dir = dir
        self.label_list = os.listdir(self.dir)
        self.path = os.path.join(self.dir)
        self.value_transform = A.Compose([
            A.Normalize(mean=[0.4970, 0.5882, 0.2305], std=[0.2359, 0.2376, 0.2131]),
            A.Resize(224, 224),
            ToTensorV2(),
        ])

    def __getitem__(self, idx):
        img_name = self.label_list[idx]
        img_item_path = os.path.join(self.dir, img_name)
        image = cv2.imread(img_item_path)
        if self.value_transform is not None:
            image = self.value_transform(image=image)['image']
        label = img_name
        return image, label

    def __len__(self):
        return len(self.label_list)

    def target_to_transform(self, label, to_dic=None):
        if to_dic is None:
            to_dic = self.train_to_dic
        return torch.tensor(to_dic[label.values[0]])

    def target_from_transform(self, label, from_dic=None):
        if from_dic is None:
            from_dic = self.train_from_dic
        return from_dic[int(label)]
if __name__ == '__main__':
    Model=torch.load("./TemporaryModel.pth", map_location='cuda')
    data=MyTestDataSet('./test_images',csv_path='./train.csv')
    dataLoader=DataLoader(data,batch_size=5)
    Model.eval()
    pre_list=[]
    for i,(da,idx) in enumerate(dataLoader):
        print(i)
        da=da.to('cuda')
        pre=Model(da)
        # print(*zip(idx,torch.max(pre.data, 1)[1]))
        pre_list.extend(zip(idx,torch.max(pre.data, 1)[1]))
    pre_df=pd.DataFrame(pre_list)
    pre_df.iloc[:,1]=pre_df.iloc[:,1].map(data.target_from_transform)
    #输出到文件
    pre_df.to_csv("pre_DenseNet121.csv",index=False,header=["image_id","label"])